#include "componentfinder.h"
#include <boost/foreach.hpp>
#include "stitcher.h"
#include "inputimage.h"
#include "imagematch.h"
#include "imageset.h"

using namespace cshotel;
using namespace std;
using namespace boost;

ComponentFinder::ComponentFinder(shared_ptr<ImageSet>  imgs) : imageSet(imgs), components(new vector<shared_ptr<ImageSet> >()) {}

bool imageSetCompareFunction(shared_ptr<ImageSet> a, shared_ptr<ImageSet> b) {
    return a->getSize() > b->getSize();
}

shared_ptr<vector<shared_ptr<ImageSet > > >ComponentFinder::findComponents() {
    dout("Finding strongly connected components in the matching graph ..." << endl);
    foreach(shared_ptr<InputImage> img, imageSet->getImages()) {
        if(hittingTime.find(img) == hittingTime.end()) {
            dfs(img);
        }
    }
    
    sort(components->begin(), components->end(), imageSetCompareFunction);
    return components;
}

void ComponentFinder::dfs(shared_ptr<InputImage> currentImage) {
    hittingTime[currentImage] = hittingTimeCounter;
    ++hittingTimeCounter;
    lowLink[currentImage] = hittingTime[currentImage];
    dfsStack.push(currentImage);
    inStack[currentImage] = true;
    
    foreach(shared_ptr<ImageMatch> match, currentImage->getImageMatches()) {
        if(shared_ptr<InputImage> next = (match->getSecondImage().lock())) {
            if(hittingTime.find(next) == hittingTime.end()) {
                dfs(next);
                lowLink[currentImage] = min(lowLink[currentImage], lowLink[next]);
            } else if(inStack[next]) {
                lowLink[currentImage] = min(lowLink[currentImage], lowLink[next]);
            }
        }
    }
    
    if(lowLink[currentImage] == hittingTime[currentImage]) {
        shared_ptr<ImageSet> newImageSet(new ImageSet());
        
        while(!dfsStack.empty()) {
            shared_ptr<InputImage> stackImage = dfsStack.top();
            dfsStack.pop();
            inStack[stackImage] = false;
            newImageSet->addImage(stackImage);
            if(stackImage == currentImage) {
                break;
            }
        }
        
        components->push_back(newImageSet);
        dout("    found new component (size " << newImageSet->getSize() << "):");
        foreach(shared_ptr<InputImage> img, newImageSet->getImages()) {
            dout(" " << img->getId());
        }
        dout(endl);
    }
}
